//      vom_main.cc
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#include <iostream>
#include <boost/program_options.hpp>
#include <boost/date_time/posix_time/posix_time.hpp> 

#include <boost/foreach.hpp> 
#include <map>
#include "stl-helper.h"

#include "dna_processing.h"
#include "suffix_array.h"
#include "vom.h"

namespace po = boost::program_options;
po::options_description desc("Allowed options");
po::variables_map vm;

void process_options(int argc, char** argv);
void usage(int exitcode = 0);

void read_fasta_test();
void suffix_array_test();
void vom_tree_test();
void fix_tree_test();
void likelihood_test();

int main(int argc, char** argv)
{
  MSG("Variable Order Model ... ");

  process_options(argc, argv);

  boost::posix_time::ptime time_start(boost::posix_time::microsec_clock::local_time());

  if (vm.count("help"))
    usage();
  else if (vm.count("read_fasta"))
    read_fasta_test();
  else if (vm.count("suffix_array"))
    suffix_array_test();
  else if (vm.count("vom_tree"))
    vom_tree_test();
  else if (vm.count("fix_tree"))
    fix_tree_test();
  else if (vm.count("likelihood"))
    likelihood_test();
  else
    usage();

  boost::posix_time::ptime time_end(boost::posix_time::microsec_clock::local_time());
  boost::posix_time::time_duration duration(time_end - time_start);
  MSG("Duration = " << duration);
}

void process_options(int argc, char** argv)
{
  desc.add_options()
    ("help", "produce this usage")
    ("input", po::value<std::string>(), "input file")
    ("output", po::value<std::string>(), "output file")
    ("model", po::value<std::string>(), "model file")
    ("seed", po::value<int>(), "random seed (REQUIRED)")

    // PARAMETERS
    ("order", po::value<int>(), "order of the model")

    //JOB TYPE
    ("read_fasta", "read fasta file\n--input=<FASTA FILE>")
    ("suffix_array", "suffix array test\n--input=<FASTA FILE>")
    ("vom_tree", "build vom tree\n--input=<FASTA FILE>")
    ("fix_tree", "build vom tree\n--input=<FASTA FILE>")
    ("likelihood", "likelihood computation\n--input=<FASTA FILE> --model=<MODEL FILE> --output=")
    ;
  po::store(po::parse_command_line(argc, argv, desc), vm);
  po::notify(vm);
}

void usage(int exitcode)
{
  std::cout << desc << std::endl;
  exit(exitcode);
}

void read_fasta_test()
{
  std::string inputFile = vm["input"].as<std::string>();
  Sequences seqs;

  MSG("Reading fasta file  ... " << inputFile);
  seqs.fasta_read(inputFile);

  if (!vm.count("output"))
    seqs.print(std::cerr);
  else {
    std::string outputFile = vm["output"].as<std::string>();
    FileAccessWrite(output, outputFile);
    seqs.print(output, true);
    output.close();
  }
}

void suffix_array_test() 
{
  std::string inputFile = vm["input"].as<std::string>();
  Sequences seqs;

  MSG("Reading fasta file  ... " << inputFile);
  seqs.fasta_read(inputFile);

  SequencesSA seqsa(seqs);

  MSG("count(ATG) = " << seqsa.count("ATG"));
  MSG("count(TTTT) = " << seqsa.count("TTTT"));
  MSG("count(AAAA) = " << seqsa.count("AAAA"));
}

void vom_tree_test()
{
  std::string inputFile = vm["input"].as<std::string>();
  Sequences seqs;

  MSG("Reading fasta file  ... " << inputFile);
  seqs.fasta_read(inputFile);

  SequencesSA seqsa(seqs);

  VOM vom;
  vom.build_big_tree(seqs, seqsa);
  vom.mle_estimate();

  if (!vm.count("output"))
    vom.print(std::cerr);
  else {
    std::string outputFile = vm["output"].as<std::string>();
    MSG("Writing tree file  ... " << outputFile);
    FileAccessWrite(output, outputFile);
    vom.print(output, true);
    output.close();
  }
}

void fix_tree_test()
{
  std::string inputFile = vm["input"].as<std::string>();
  int order = vm["order"].as<int>();
  Sequences seqs;

  MSG("Reading fasta file  ... " << inputFile);
  seqs.fasta_read(inputFile);

  SequencesSA seqsa(seqs);

  VOM vom;
  vom.build_fix_tree(seqs, seqsa, order);
  vom.mle_estimate();

  if (!vm.count("output"))
    vom.print(std::cerr);
  else {
    std::string outputFile = vm["output"].as<std::string>();
    MSG("Writing tree file  ... " << outputFile);
    FileAccessWrite(output, outputFile);
    vom.print(output, true);
    output.close();
  }
}

void likelihood_test()
{
  std::string inputFile = vm["input"].as<std::string>();
  std::string modelFile = vm["model"].as<std::string>();
  std::string outputFile = vm["output"].as<std::string>();
  int order = vm.count("order") ? vm["order"].as<int>() : -1;

  Sequences seqs;

  MSG("Reading fasta file  ... " << inputFile);
  seqs.fasta_read(inputFile);

  VOM vom;
  MSG("Reading model file  ... " << modelFile);
  vom.load(modelFile);

  /*
  std::string outputFile = modelFile+".test";
  FileAccessWrite(output, outputFile);
  vom.print(output, true);
  output.close();
  */

  MSG("Writing log-likelihood to output file  ... " << outputFile);
  FileAccessWrite(output, outputFile);
  for (uint i = 0; i < seqs.size(); i++) {
    double ll = vom.log_likelihood(seqs[i], order);
    output << ">" << seqs.annotation_[i] << ENDL << ll << ENDL;
  }
  output.close();
}
